import os
import hashlib
from glob import glob
from tqdm import tqdm
import av
from PIL import Image
import sys
try:
    from DWPose.inference import get_pose
except ImportError:
    # 如果未找到模块，则临时添加上层目录
    sys.path.append("../..")
    from DWPose.inference import get_pose
    sys_path_inserted = True


#=========================================工具函数========================================

def get_video_hash(video_path, length=8):
    """为视频文件生成一个简短的、唯一的哈希值。"""
    sha = hashlib.sha256()
    with open(video_path, "rb") as f:
        while chunk := f.read(8192):
            sha.update(chunk)
    return sha.hexdigest()[:length]

#=========================================核心处理模块========================================

def extract_frames(video_path, frame_output_folder):
    """
    从视频中提取关键帧。
    """
    os.makedirs(frame_output_folder, exist_ok=True)
    print(f" extracting frames from: {video_path}")
    
    try:
        with av.open(video_path) as container:
            stream = container.streams.video[0]
            # 设置只解码关键帧以加快速度
            stream.codec_context.skip_frame = "NONKEY"

            for i, frame in enumerate(container.decode(stream)):
                output_filename = os.path.join(frame_output_folder, f"frame_{i:05d}.png")
                frame.to_image().save(output_filename)
        
        print(f"✅ 抽帧完成 -> {frame_output_folder}")
    except Exception as e:
        print(f"❌ 抽帧失败: {video_path}, 错误: {e}")


def run_pose_detection_on_folder(input_image_dir, output_pose_dir):
    """
    对指定文件夹中的所有图片进行姿态识别。
    这是一个通用函数，可用于处理视频帧或图片目录。
    """
    success_cnt = 0
    failed_cnt = 0
    os.makedirs(output_pose_dir, exist_ok=True)
    
    # 支持多种常见图片格式
    image_paths = []
    for ext in ("*.png", "*.jpg", "*.jpeg", "*.bmp"):
        image_paths.extend(glob(os.path.join(input_image_dir, ext)))

    if not image_paths:
        print(f"⚠️ 在目录 {input_image_dir} 中未找到任何图片，跳过 Pose 检测。")
        return

    for img_path in tqdm(image_paths, desc=f"🏃‍ Pose 识别于 {os.path.basename(input_image_dir)}"):
        try:
            pose_img = get_pose(img_path)  # 返回 PIL.Image 对象
            if pose_img is None:
                failed_cnt+=1
                # 如果存在，删除
                output_path = os.path.join(output_pose_dir, os.path.basename(img_path))
                if os.path.exists(output_path):
                    os.remove(output_path)
                    print(f"⚠️ 删除已存在的错误文件: {output_path}")
                print(f"- 跳过无法处理的图片: {img_path}")
                with open("no_pose.txt", "a") as f:
                    f.write(img_path + "\n")
                continue
            
            # 保持原始文件名
            output_path = os.path.join(output_pose_dir, os.path.basename(img_path))
            pose_img.save(output_path)
            success_cnt += 1
        except Exception as e:
            print(f"❌ 处理失败: {img_path}，错误: {e}")
            with open('error.log', 'a') as f:
                f.write(f"{img_path}\n")

    print(f"✅ Pose 检测完成 -> {output_pose_dir} (成功: {success_cnt}, 失败: {failed_cnt})")
    with open("no_pose.txt", "a") as f:
        f.write("==============\n")


#=========================================主流程控制器========================================

def process_video_directory(root_video_dir, output_base_dir, skip_exist=True):
    """
    主流程：处理指定目录下的所有视频文件。
    """
    print("--- 模式: 视频处理 ---")
    video_extensions = ('.mp4', '.MP4')

    for path, _, files in os.walk(root_video_dir):
        for file in sorted(files):
            if file.endswith(video_extensions):
                video_path = os.path.join(path, file)
                video_hash = get_video_hash(video_path)
                print(f"\n🎥 处理视频: {file} (hash: {video_hash})")

                frame_output_folder = os.path.join(output_base_dir, "frames", video_hash)
                pose_output_folder = os.path.join(output_base_dir, "poses", video_hash)

                # 1. 抽帧
                if os.path.exists(frame_output_folder) and skip_exist:
                    print(f"🔁 已存在抽帧结果，跳过: {frame_output_folder}")
                else:
                    extract_frames(video_path, frame_output_folder)

                # 2. 姿态识别
                if os.path.exists(pose_output_folder) and skip_exist:
                    print(f"🔁 已存在 Pose 数据，跳过: {pose_output_folder}")
                else:
                    run_pose_detection_on_folder(frame_output_folder, pose_output_folder)


def process_image_directory(root_image_dir, output_base_dir, skip_exist=True):
    """
    主流程：处理指定目录下的所有图片子目录。
    """
    print("--- 模式: 图片处理 ---")
    
    # 查找所有一级子目录 (e.g., class_A, class_B)
    # os.scandir 更高效
    subdirs = [d.path for d in os.scandir(root_image_dir) if d.is_dir()]

    if not subdirs:
        print(f"❌ 在 '{root_image_dir}' 下未找到任何子目录。请确保图片按子目录存放。")
        return

    for image_group_dir in sorted(subdirs):
        group_name = os.path.basename(image_group_dir)
        print(f"\n🖼️ 处理图片组: {group_name}")

        # 直接将图片组作为姿态识别的输入
        # 输出目录也使用组名来保持结构一致
        pose_output_folder = os.path.join(output_base_dir, "poses", group_name)
        
        # 姿态识别
        if os.path.exists(pose_output_folder) and skip_exist:
            print(f"🔁 已存在 Pose 数据，跳过: {pose_output_folder}")
        else:
            # 输入是图片组的目录，输出是对应的pose目录
            run_pose_detection_on_folder(image_group_dir, pose_output_folder)


#=========================================程序入口========================================

if __name__ == "__main__":
    
    # --- 请在这里配置 ---
    
    # 1. 选择处理类型: 'video' 或 'image'
    INPUT_TYPE = 'image'  # <-- 修改这里来切换模式

    # 2. 设置输入路径
    #    - 如果 INPUT_TYPE 是 'video', 这里应该是存放视频的根目录 (e.g., "videos")
    #    - 如果 INPUT_TYPE 是 'image', 这里应该是存放图片子目录的根目录 (e.g., "image_resized")
    INPUT_PATH = "image_resized" # <-- 修改这里来指定输入源

    # 3. 设置总的输出目录
    OUTPUT_BASE_DIR = "output/"

    # 4. 是否跳过已存在的处理结果
    SKIP_IF_EXIST = False

    # --- 配置结束 ---

    # 根据配置调用相应的处理函数
    if INPUT_TYPE == 'video':
        process_video_directory(INPUT_PATH, OUTPUT_BASE_DIR, skip_exist=SKIP_IF_EXIST)
    elif INPUT_TYPE == 'image':
        process_image_directory(INPUT_PATH, OUTPUT_BASE_DIR, skip_exist=SKIP_IF_EXIST)
    else:
        print(f"❌ 错误: 无效的 INPUT_TYPE '{INPUT_TYPE}'。请选择 'video' 或 'image'。")

    print("\n🎉 所有任务处理完成！")